package com.coder.algorithm.leetcode

/**
 * [498. 对角线遍历](https://leetcode-cn.com/problems/diagonal-traverse/)
 *
 * @author yuhus
 * @date: 2022/3/14
 */
class l_0498 {
    /**
     * 规律：
     * 1、元素共有 M*N 个
     * 2、此种遍历，共有 N+M-1 行
     * 3、向上起始位置：（第一列，最后一行）结束位置：（第一行，最后一列）
     * 3、向下起始位置：（第一行，最后一列）结束位置：（第一列，最后一行）
     */
    fun findDiagonalOrder(mat: Array<IntArray>): IntArray {
        val M = mat.size // 行
        val N = mat[0].size// 列
        val numbers = IntArray(M * N)
        var i = 0
        // 可以按照行进行遍历（方向和行相关），也可以按照元素个数遍历
        for (row in 0 until N + M -1) {
            // 向上
            val directionUR = row % 2 == 0
            // 起始位置，第一列/最后一行（UR方向，反(DL)方向起始位置和终止位置互换）
            var startRow: Int
            var startCol: Int
            if (row < M) {// 第一列
                startRow = row
                startCol = 0
            } else { // 最后一行
                startRow = M - 1
                startCol = row - M + 1
            }
            // 结束位置
            var endRow: Int
            var endCol: Int
            if (row < N) { // 第一行
                endRow = 0
                endCol = row
            } else { // 最后列
                endRow = row - N + 1
                endCol = N - 1
            }
            if (directionUR) {
                var r = startRow
                var c = startCol
                while (r >= 0 && c < N) {
                    numbers[i++] = mat[r][c]
                    r -= 1
                    c += 1
                }
            } else {
                var r = endRow
                var c = endCol
                while (r < M && c >= 0) {
                    numbers[i++] = mat[r][c]
                    r += 1
                    c -= 1
                }
            }
        }
        return numbers
    }

    fun findDiagonalOrder_reduce(mat: Array<IntArray>): IntArray {
        val M = mat.size // 行
        val N = mat[0].size// 列
        val numbers = IntArray(M * N)
        var i = 0
        // 可以按照行进行遍历（方向和行相关），也可以按照元素个数遍历
        for (row in 0 until N + M -1) {
            // 向上
            val directionUR = row % 2 == 0
            var r = if (directionUR) {
                if (row < M) { row } else { M - 1 }
            } else {
                if (row < N) { 0 } else { row - N + 1}
            }
            var c = if (directionUR) {
                if (row < M) { 0 } else { row - M + 1 }
            } else {
                if (row < N) { row } else { N - 1 }
            }
            while (r in 0 until M && c in 0 until N) {
                numbers[i++] = mat[r][c]
                r = if (directionUR) { r - 1} else { r + 1 }
                c = if (directionUR) { c + 1} else { c - 1 }
            }
        }
        return numbers
    }
}